import numpy as np
import math as math
import matplotlib.pyplot as plt
import os, sys

#filename = ['/home/kjs/H2O_cross_sxns/H2O_prescriptions.txt', 
#            '/home/kjs/gadams_sample_mars_one/h2o_cross_section/CROSS_H2O=H+OH_LORES1.DAT_arthur',
#            '/home/kjs/gadams_sample_mars_one/h2o_cross_section/CROSS_H2O=H+OH_LORES1.DAT_frank']

# dja replacing the filenames - want H2O=H2O for cross sections & not quantum yields 
# for ranjan, using .DAT format for kinetics instead of the original txt data (so formatting matches)
filename = ['/home/kjs/gadams_sample_mars_one/Cross_arthur/CROSS_H2O=H2O_LORES1.DAT',\
            '/home/kjs/gadams_sample_mars_one/Cross_arthur/CROSS_H2O=H2O_LORES1.DAT',\
            '/home/kjs/gadams_sample_mars_one/Cross_arthur_duplicate/CROSS_H2O=H2O_LORES1.DAT'] # frank
            
fctr = 0
def ifile(filename):
    infile = open(filename, 'r')
    temp = infile.readline() # dja just tabbing things in
    for iskip in range(0,3): # dja skip 4 lines of header instead of 2                                                                                                                                     
        temp = infile.readline() # dja skip 4 lines of header instead of 2

    wvl, extrap_pre, cutoff_pre  = [],[],[]
    temp = infile.readline().split()
    while len(temp)>1:
        print(temp)
        wvl += [float(temp[0])] # wavelength is in the 1st column of data; so use index 0
        extrap_pre  += [float(temp[1])] # extrapolation prescription is in the 2nd column of data; so use index 1
# dja remove these after replacing above file with .DAT (then all should have same format)
#    if len(temp)>2:
#        cutoff_pre += [float(temp[2])] # cutoff prescription is in the 3rd column of data; so we use index 2
#    else:
#        cutoff_pre+=[float('nan')]
        temp = infile.readline().split() # read next line and separate with a space                                                             
    dat = {'wvl':wvl,'cross':extrap_pre}
    return dat

                                                       
#print(dat['wvl'], dat['cross']) # print values to screen (can check in txt file to ensure values match)
fig,ax=plt.subplots(figsize=(8,8)) # open a blank document                                                                                      

clrs = ['black','blue','red']
lns = ['-','--','--']
cases = ['ranjan','arthur','frank']

for ifilename in filename:
    dat= ifile(ifilename) # read one file
    if not cases[fctr] == 'ranjan':
        ax.semilogy(dat['wvl'][0:92], dat['cross'][0:92], color = clrs[fctr], linewidth=2.5, linestyle=lns[fctr], label=cases[fctr])
    else:
        ax.semilogy(dat['wvl'][91:], dat['cross'][91:], color = clrs[fctr], linewidth=2.5, linestyle=lns[fctr], label=cases[fctr])
    fctr+=1

#ax.semilogy(wvl, cutoff_pre, color = 'red', alinewidth=2.5, linestyle='-', label='cutoff_pre')

ax.set_xlabel('wvl')
ax.set_ylabel('cross sections')
ax.legend() # comment: a small box inside your figure will label the two lines
plt.savefig('cross_sections.png')
exit()

# ignore below

#ax.set_xlabel('wvl')
#ax.set_ylabel('H2O_pre')

# reading in the headers for the CROSS_H2O=H+OH_LORES1.DAT_arthur file, there's 4 lines before data starts
for iskip in range(0,4):
    temp = infile.readline()
temp = infile.readline().split()
# keep reading until no more data                                                                                                                                                                         
while len(temp)>1:
    print(temp)
wvl_1, cross_sections = [],[]
temp = infile.readline().split()
# keep reading until no more data                                                                                                                                                                          
wvl_1 += [float(temp[0])] # wavelength is in the 1st column of data; so use index 0                                                                                                                    
cross_section  += [float(temp[1])] # cross section  is in the 2nd column of data; so use index 1                                                                                                       
print(wvl_1, cross_section)

# reading in the headers for the CROSS_H2O=H+OH_LORES1.DAT_frank file, there's 4 lines before data starts
for iskip in range(0,4):
    temp = infile.readline()
wvl_2, cross_sections = [],[]
temp = infile.readline().split()
# keep reading until no more data                                                                                                                                                                      
while len(temp)>1:
    print(temp)
    wvl_2 += [float(temp[0])] # wavelength is in the 1st column of data; so use index 0                                                                                                                    
    cross_section  += [float(temp[1])] # cross section  is in the 2nd column of data; so use index 1  
temp = infile.readline().split()
print(wvl_2, cross_sections)

fig,ax=plt.subplots(figsize=(8,8)) # open a blank document
ax.semilogy(wvl, cross_sections, color = 'blue', linewidth=2.5, linestyle='-', label='cross_sxns_ranjan')
ax.semilogy(wvl, cross_sections, color = 'red', linewidth=2.5, linestyle='--', label='cross_sxns_arthur')
ax.semilogy(wvl, cross_sections, color = 'cyan', linewidth=2.5, linestyle='-.', label='cross_sxns_frank')
ax.set_xlabel('wvl')
ax.set_ylabel('cross sections')
ax.legend()

outname = 'ranjan_arthur_frank.png'
plt.savefig(outname)
